
import cv2
import numpy as np
import matplotlib
matplotlib.use('TkAgg') 

min_length = 9999 

def extract_white_paper(img):
    # 转换到HSV颜色空间（便于分离亮度和饱和度）
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)
    
    # 白纸特征：低饱和度（s）+ 高亮度（v）
    # 阈值可根据实际环境调整（此处为示例值）
    s_threshold = 60  # 饱和度阈值（低饱和度）
    v_threshold = 120  # 亮度阈值（高亮度）
    s_mask = s < s_threshold
    v_mask = v > v_threshold
    white_mask = np.logical_and(s_mask, v_mask).astype(np.uint8) * 255  # 合并掩码
    
    # 形态学操作：去除小噪声，填充孔洞
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    white_mask = cv2.morphologyEx(white_mask, cv2.MORPH_CLOSE, kernel, iterations=3)  # 填充孔洞
    white_mask = cv2.morphologyEx(white_mask, cv2.MORPH_OPEN, kernel, iterations=1)   # 去除小噪声
    
    # 查找轮廓，筛选最大的连通区域（假设白纸是最大的白色区域）
    contours, _ = cv2.findContours(white_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None, None  # 未找到白纸区域
    max_contour = max(contours, key=cv2.contourArea)
    area = cv2.contourArea(max_contour)
    if area < 10000:  # 过滤过小的区域（避免误判）
        return None, None
    
    # 创建白纸区域的掩码（仅保留最大轮廓内的区域）
    paper_mask = np.zeros_like(white_mask)
    cv2.drawContours(paper_mask, [max_contour], -1, 255, -1)
    
    # 提取白纸区域（原图中仅保留掩码内的像素）
    paper_img = cv2.bitwise_and(img, img, mask=paper_mask)
    return paper_img, paper_mask

def extract_lines(paper_img, paper_mask):
    # 转换为灰度图
    gray = cv2.cvtColor(paper_img, cv2.COLOR_BGR2GRAY)
    
    # 反色处理（让线条变为白色，背景为黑色）
    # 先通过掩码去除白纸外的干扰（置为黑色）
    gray[paper_mask == 0] = 0
    
    inverted = cv2.bitwise_not(gray)  # 反色：白纸（亮）→ 黑，线条（暗）→ 白
    # 通过掩码去除白纸外的干扰（置为黑色）
    inverted[paper_mask == 0] = 0
    # 阈值处理：保留线条（白色）
    _, line_mask = cv2.threshold(inverted, 160, 255, cv2.THRESH_BINARY) 
    # 形态学优化：连接断裂的线条，去除小噪声
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    line_mask = cv2.morphologyEx(line_mask, cv2.MORPH_CLOSE, kernel, iterations=2)  # 连接线条
    line_mask = cv2.morphologyEx(line_mask, cv2.MORPH_OPEN, kernel, iterations=1)   # 去噪声
    return line_mask

def prune_short_branches(skeleton, min_length):
    img = skeleton.copy()
    _, binary = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
    pruned = binary.copy()
    visited = np.zeros_like(binary, dtype=bool)
    devisited = np.zeros_like(binary, dtype=bool)
    # 8邻域方向
    neighbors = [(-1, -1), (-1, 0), (-1, 1),
                 (0, -1),          (0, 1),
                 (1, -1),  (1, 0), (1, 1)]

    def is_endpoint(x, y,img):
        if img[x, y] != 255:
            return False
        count = 0
        for dx, dy in neighbors:
            nx, ny = x + dx, y + dy
            if 0 <= nx < img.shape[0] and 0 <= ny < img.shape[1]:
                if img[nx, ny] == 255:
                    count += 1
        return count == 1
    
    def dfs_branch_iterative(start_x, start_y,img):
        # 栈元素：(当前x, 当前y, 上一x, 上一y)
        stack = [(start_x, start_y, -1, -1)]
        path = []
        visited[start_x, start_y] = True  # 标记起点为已访问
        
        while stack:
            x, y, prev_x, prev_y = stack.pop()  # 弹出栈顶元素处理
            path.append((x, y))  # 记录当前像素到路径
            
            # 寻找下一个有效像素（排除上一个像素）
            next_pixels = []
            for dx, dy in neighbors:
                nx, ny = x + dx, y + dy
                if 0 <= nx < img.shape[0] and 0 <= ny < img.shape[1]:
                    # 条件：前景像素、未访问、不是上一个像素
                    if img[nx, ny] == 255 and ((not visited[nx, ny]) or devisited[nx, ny]) and (nx != prev_x or ny != prev_y):
                        next_pixels.append((nx, ny))
            
            # 分支点：若有多个下一跳，当前分支结束（只沿一个方向继续遍历）
            if len(next_pixels) == 1:
                nx, ny = next_pixels[0]
                visited[nx, ny] = True  # 提前标记为已访问，避免重复处理
                stack.append((nx, ny, x, y))  # 下一个像素入栈
            # 若有0个下一跳（终点）或≥2个下一跳（分支点），则当前分支结束
            else:
                break
        
        return path
    
    # 遍历图像，检测所有端点并处理

    for x in range(binary.shape[0]):
        for y in range(binary.shape[1]):
        # 仅处理未访问的端点
            if binary[x, y] == 255 and ((not visited[x, y]) or devisited[x, y]) and is_endpoint(x, y,pruned):
            # 用迭代版DFS替代递归，避免深度问题
                branch_path = dfs_branch_iterative(x, y,pruned)
                branch_length = len(branch_path)

            # 显示中间结果（可根据需要保留或删除）
                   
            # 若分支长度小于阈值，则删除该分支（忽略最后一个元素）
                if branch_length < min_length:
                    devisited[branch_path[-1]] = True
                    for (px, py) in branch_path[:-1]:  # 忽略最后一个元素
                        pruned[px, py] = 0

    
    return pruned
  
def skeletonized(line_mask):
    skeleton = cv2.ximgproc.thinning(line_mask,thinningType=cv2.ximgproc.THINNING_GUOHALL)
    skeleton = prune_short_branches(skeleton,min_length)
    return skeleton

def detect_circles(gray_img, dp=1, min_dist=20, param1=50, param2=30, min_radius=0, max_radius=0):
    
    # 检查输入是否为单通道图像
    if len(gray_img.shape) != 2:
        raise ValueError("输入必须是单通道（灰度）图像！")
    
    # 预处理：高斯模糊去噪（减少噪声对圆检测的干扰）
    blurred = cv2.GaussianBlur(gray_img, (5, 5), 0)  # 5x5卷积核， sigma=0（自动计算）
    
    # 霍夫圆变换检测圆（使用梯度法）
    circles = cv2.HoughCircles(
        image=blurred,
        method=cv2.HOUGH_GRADIENT,
        dp=dp,
        minDist=min_dist,
        param1=param1,
        param2=param2,
        minRadius=min_radius,
        maxRadius=max_radius
    )
    
    # 提取圆心坐标（转换为整数坐标）
    centers = []
    if circles is not None:
        # 霍夫变换返回的circles是形状为(1, N, 3)的数组，N为圆的数量，每个圆为(x, y, r)
        circles = np.uint16(np.around(circles))  # 四舍五入并转为整数
        for circle in circles[0, :]:
            x, y, _ = circle  # 只取圆心坐标(x, y)，忽略半径r
            centers.append([x,y])

    return centers

def fit_shape_with_douglas_peucker(foreground_img, line_mask, epsilon=5.0):
    # 确保输入图像是单通道的
    if len(foreground_img.shape) == 3:
        gray = cv2.cvtColor(foreground_img, cv2.COLOR_BGR2GRAY)
    else:
        gray = foreground_img.copy()
    
    # 确保图像是二值化的
    if gray.max() > 1:
        _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
    else:
        binary = (gray * 255).astype(np.uint8)
    
    # 查找轮廓
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    fitted_contours = []
    # 对每个轮廓应用Douglas-Peucker算法
    for contour in contours:
        # 只处理有足够点的轮廓
        if len(contour) > 3:
            # 使用approxPolyDP进行拟合（Douglas-Peucker算法）
            approx = cv2.approxPolyDP(contour, epsilon, closed=True)
            # 将点集转换为更易处理的格式
            fitted_points = [point[0] for point in approx]
            fitted_contours.append(fitted_points)
    
    return fitted_contours

def find_nearest_point(points, target_point):
    # 将点列转换为NumPy数组（兼容list和数组输入）
    points_np = np.array(points, dtype=np.float32)
    
    # 处理伪三维数组（形状为(N,1,2)）：去除冗余的中间维度（1）
    if points_np.ndim == 3:
        # 检查是否为 (N,1,2) 结构（伪三维的典型特征）
        if points_np.shape[1] == 1 and points_np.shape[2] == 2:
            points_np = points_np.reshape(-1, 2)  # 转换为 (N,2)
        else:
            raise ValueError("三维点列格式错误！需为伪三维数组（如[[[x1,y1]], [[x2,y2]], ...]，形状(N,1,2)）")
    
    # 验证转换后的数组是否为 (N,2) 格式
    if points_np.ndim != 2 or points_np.shape[1] != 2:
        raise ValueError("点列格式错误！支持格式：二维list([[x1,y1],...])、伪三维数组([[[x1,y1]],...])、(N,2)数组")
    
    # 处理目标点（统一为(1,2)形状
    target_np = np.array(target_point, dtype=np.float32).reshape(1, 2)
    
    # 5. 计算距离平方并找到最近点
    distances_sq = np.sum((points_np - target_np) **2, axis=1)
    min_idx = np.argmin(distances_sq)
    nearest_point = tuple(map(int, points_np[min_idx]))
    
    return nearest_point

def process_image(img):
    # 读取图像
    if img is None:
        print("无法读取帧，请检查")
        return
    # 保存原图（可选）
    #show_img("原图", img, save_path="OutPut2/original.png")
    
    # 步骤1：提取白纸区域
    paper_img, paper_mask = extract_white_paper(img)
    if paper_img is None:
        print("未检测到白纸区域")
        return

    line_mask = extract_lines(paper_img, paper_mask)

    skeleton = skeletonized(line_mask)

    fitted_contours = fit_shape_with_douglas_peucker(skeleton, line_mask, epsilon=4)
    annotated_img = img.copy()
    center_list = detect_circles(line_mask,min_radius=10,max_radius=100)
    for i, contour in enumerate(fitted_contours):
        # 绘制拟合线
        points = np.array(contour, np.int32)
        points = points.reshape((-1, 1, 2))
        cv2.polylines(annotated_img, [points], True, (0, 0, 255), 2) # 红色线条
        # 绘制拟合点
        for j, (x, y) in enumerate(contour):
            cv2.circle(annotated_img, (x, y), 5, (0, 255, 0), -1)  # 绿色点
            cv2.putText(annotated_img, f"{j}", (x+10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)  # 点编号
        if len(center_list)>i:
            cv2.circle(annotated_img, center_list[i], 5, (0, 255, 0), -1)  # 绿色点
            first_point=find_nearest_point(points, center_list[i])
            cv2.line(annotated_img, first_point,center_list[i], (0, 0, 255), 2)
        # 绘制起点与图形连线
    return annotated_img#,line_mask,paper_img

def process_camera_video():
    # 打开默认摄像头（参数0表示默认摄像头，多摄像头时可尝试1、2等）
    cap = cv2.VideoCapture(0)
    
    # 检查摄像头是否成功打开
    if not cap.isOpened():
        print("无法打开摄像头，请检查设备是否连接或权限是否足够")
        return
    
    # 创建显示窗口
    cv2.namedWindow("原始图像", cv2.WINDOW_NORMAL)
    cv2.namedWindow("处理后图像", cv2.WINDOW_NORMAL)
    #cv2.namedWindow("line_mask", cv2.WINDOW_NORMAL)
    #cv2.namedWindow("paper_img", cv2.WINDOW_NORMAL)

    try:
        while True:
            # 读取一帧图像
            # ret：布尔值，表示是否成功读取帧
            # frame：读取到的帧（BGR格式）
            ret, frame = cap.read()
            
            # 如果读取失败，退出循环
            if not ret:
                print("无法获取图像帧，可能摄像头已断开")
                break
            
            # 调用你的图片处理函数处理当前帧
            processed_frame = process_image(frame)
            
            # 显示原始图像和处理后的图像
            cv2.imshow("原始图像", frame)
            cv2.imshow("处理后图像", processed_frame)
            #cv2.imshow("line_mask", line_mask)
            #cv2.imshow("paper_img", paper_img)

            # 等待1毫秒，检测按键（按q键退出）
            key = cv2.waitKey(1)
            if key == ord('q') or key == 27:  # 27是ESC键
                print("退出程序")
                break
    finally:
        # 释放摄像头资源
        cap.release()
        # 关闭所有显示窗口
        cv2.destroyAllWindows()

if __name__ == "__main__":
    process_camera_video()